import tensorflow as tf
import os
from util import *


class CycleGAN():

    def __init__(self, generatorG, generatorF, disciminatorDY, discriminatorX,
                 genOptimizer, discOptimizer, logDir = "./log"):
        # self.inputShape =
        # self.disciminator = discriminator
        # self.generator = generator
        # self.mode = mode
        self.genOptimizer = genOptimizer
        self.discOptimizer = discOptimizer
        self.buildModel(generatorG, generatorF, disciminatorDY, discriminatorX)

    def buildModel(self, generatorG, generatorF, disciminatorDY, discriminatorX):
        self.G = generatorG
        self.F = generatorF
        self.Dy = disciminatorDY
        self.Dx = discriminatorX

    def computeLoss(self, x, y, lambdaCyc, lambdaId):
        y_ = self.G(x)
        x_ = self.F(y_)
        Fx = self.F(x)
        Gy = self.G(y)
        DxReal = self.Dx(x)
        DxFake = self.Dx(x_)
        DyReal = self.Dy(y)
        DyFake = self.Dy(y_)
        cycDistanceXtoFakeX = distancel1(x, x_)
        cycDistanceYtoFakeY = distancel1(y, y_)
        idDistanceXtoX = distancel1(x, Fx)
        idDistanceYtoY = distancel1(y, Gy)
        cycLoss = cycDistanceYtoFakeY + cycDistanceXtoFakeX
        idLoss = idDistanceYtoY + idDistanceXtoX
        discDyReal = lossCompute(DyReal, True)
        discDyFake = lossCompute(DyFake, False)
        discDxReal = lossCompute(DxReal, True)
        discDxFake = lossCompute(DxFake, False)
        discDyFake = tf.cast(discDyFake, tf.float64)
        discDxFake = tf.cast(discDxFake, tf.float64)
        discDxReal = tf.cast(discDxReal, tf.float64)
        discDyReal = tf.cast(discDyReal, tf.float64)
        discDy = discDyReal + discDyFake
        discDx = discDxReal + discDxFake
        genDy = lossCompute(DxFake, True)
        genDx = lossCompute(DyFake, True)
        genDx = tf.cast(genDx, tf.float64)
        genDy = tf.cast(genDy, tf.float64)
        cycLoss = tf.cast(cycLoss, tf.float64)
        idLoss = tf.cast(idLoss, tf.float64)
        genLoss = genDx + genDy + lambdaCyc * cycLoss + lambdaId * idLoss
        discLoss = discDx + discDy
        return genLoss, discLoss

    def applyGradients(self, x, y, lambdaCyc, lambdaId):
        with tf.GradientTape() as genGTape, tf.GradientTape() as discXTape, tf.GradientTape() as genFTape, tf.GradientTape() as discYTape:
            genLoss, discLoss = self.computeLoss(x, y, lambdaCyc, lambdaId)
        dxGradients = discXTape.gradient(discLoss, self.Dx.trainable_variables)
        dyGradients = discYTape.gradient(discLoss, self.Dy.trainable_variables)
        gGradients = genGTape.gradient(genLoss, self.G.trainable_variables)
        fGradients = genFTape.gradient(genLoss, self.F.trainable_variables)
        self.genOptimizer.apply_gradients(zip(gGradients, self.G.trainable_variables))
        self.genOptimizer.apply_gradients(zip(fGradients, self.F.trainable_variables))
        self.discOptimizer.apply_gradients(zip(dxGradients, self.Dx.trainable_variables))
        self.discOptimizer.apply_gradients(zip(dyGradients, self.Dy.trainable_variables))
        return genLoss, discLoss

    def train(self, x, y, lambdaCyc, lambdaId):
        genLoss, discLoss = self.applyGradients(x, y, lambdaCyc, lambdaId)
        return genLoss, discLoss

    def generate(self, inputs, direction):
        inputsTensor = tf.convert_to_tensor(inputs)
        if direction == "A2B":
            return self.G(inputsTensor)
        elif direction == "B2A":
            return self.F(inputsTensor)
        else:
            raise Exception("DirectionError.")

    def saveCkpt(self, name, dir = "./training_checkpoints"):
        checkpoint_dir = dir
        checkpoint_prefix = os.path.join(checkpoint_dir, name)
        checkpoint = tf.train.Checkpoint(generator_optimizer=self.genOptimizer,
                                         discriminator_optimize=self.discOptimizer,
                                         G=self.G,
                                         F=self.F,
                                         Dy=self.Dy,
                                         Dx=self.Dx)
        checkpoint.save(file_prefix=checkpoint_prefix)

    def loadCkpt(self, dir = "./training_checkpoints"):
        checkpoint_dir = dir
        checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
        checkpoint = tf.train.Checkpoint(generator_optimizer=self.genOptimizer,
                                         discriminator_optimize=self.discOptimizer,
                                         G=self.G,
                                         F=self.F,
                                         Dy=self.Dy,
                                         Dx=self.Dx)
        checkpoint.restore(checkpoint_prefix).assert_consumed()

    def summary(self):
        return self.G.summary(), self.F.summary(), self.Dx.summary(), self.Dy.summary()